// ProbePCAVisualizer.cs
// MonoBehaviour for visualizing original vs PCA-decoded Light Probes in scene
// Supports average error display, sampling direction, and probe selection

using System;
using System.Text;
using UnityEngine;
using UnityEngine.Rendering;

[ExecuteAlways]
public class GPUProbePCAVisualizer : MonoBehaviour
{
    public LightProbeGroup lightProbeGroup;
    public int lightProbeCountInGroup = 10;
    public GameObject visualGOTemp;
    public GameObject visualGORoot;
    
    [Header("PCA Compressed Data")]
    public PCAUtils.PrincipalAxisMethod pcaMethod = PCAUtils.PrincipalAxisMethod.Jacobi;
    public int pcaCmponents = 4;
    public CompressedProbeDataGPU compressedProbeDataGPU;
    public CompressedProbeData compressedProbeData;

    public enum SampleDirection {MainLight, SceneCamera, Up, Down, Right, Left }
    public SampleDirection sampleDirection = SampleDirection.Up;

    //public enum Channel { RGB, R, G, B }
    public ProbePCAVisualizer.Channel channel = ProbePCAVisualizer.Channel.RGB;

    [HideInInspector]
    public int selectedProbe = -1;

    public SphericalHarmonicsL2[] originalSHs;
    public SphericalHarmonicsL2[] decodedSHs;
    public Vector3[] positions;
    public float[] errors;

    public static float CalculateTotalAverageError(SphericalHarmonicsL2[] originalSHs, SphericalHarmonicsL2[] decodedSHs)
    {
        if (originalSHs.Length != decodedSHs.Length)
        {
            throw new ArgumentException("The number of SH data should be the same for original and decoded.");
        }

        float totalError = 0;

        // 遍历所有数据集，计算每个通道的误差并计算总误差
        for (int i = 0; i < originalSHs.Length; i++)
        {
            SphericalHarmonicsL2 original = originalSHs[i];
            SphericalHarmonicsL2 decoded = decodedSHs[i];

            // 计算误差
            float error = CalculateChannelError(original, decoded);

            // 将误差累加
            totalError += error;
        }

        // 返回总的平均误差
        return totalError / originalSHs.Length;
    }
    
    // 计算 Spherical Harmonics L2 的误差
    private static float CalculateChannelError(SphericalHarmonicsL2 originalSH, SphericalHarmonicsL2 decodedSH)
    {
        float totalError = 0;

        // 遍历阶数 (L) 和分量 (m)，计算每个系数的误差
        for (int L = 0; L < 3; L++) // L阶数: 0, 1, 2
        {
            for (int m = 0; m < 9; m++) // m的范围是 -L 到 L
            {
                totalError += Mathf.Abs(originalSH[L, m] - decodedSH[L, m]);
            }
        }

        // 返回平均误差
        return totalError / 9;  // 总共9个系数 (L0, L1, L2 阶数的系数)
    }
    
    public static float ComputeSHError(SphericalHarmonicsL2 original, SphericalHarmonicsL2 decoded)
    {
        float error = 0f;

        for (int channel = 0; channel < 3; channel++) // R G B
        {
            for (int coeff = 0; coeff < 9; coeff++) // L0-L8
            {
                float diff = original[channel, coeff] - decoded[channel, coeff];
                error += diff * diff;
            }
        }

        return Mathf.Sqrt(error); // 最后开平方，得到 L2 误差
    }
    
    public static float ComputeError(Color original, Color decoded)
    {
        return Vector3.Distance(
            new Vector3(original.r, original.g, original.b),
            new Vector3(decoded.r, decoded.g, decoded.b));
        
        //float error = Mathf.Sqrt(Mathf.Pow(original.r - decoded.r, 2f) +
        //                         Mathf.Pow(original.g - decoded.g, 2f) +
        //                         Mathf.Pow(original.b - decoded.b, 2f));
    }
    

    public static Color FilterChannel(Color c, ProbePCAVisualizer.Channel ch = ProbePCAVisualizer.Channel.RGB)
    {
        switch (ch)
        {
            case ProbePCAVisualizer.Channel.R:   return new Color(c.r, 0,     0);
            case ProbePCAVisualizer.Channel.G:   return new Color(0,     c.g,   0);
            case ProbePCAVisualizer.Channel.B:   return new Color(0,     0,     c.b);
            default:          return c;
        }
    }

    public Vector3 GetSampleDirection()
    {
        switch (this.sampleDirection)
        {
            case SampleDirection.MainLight:
                if (RenderSettings.sun != null)
                    return -RenderSettings.sun.transform.forward;
                break;
#if UNITY_EDITOR
            case SampleDirection.SceneCamera:
                return UnityEditor.SceneView.lastActiveSceneView.camera.transform.forward;
                break; 
#endif
            case SampleDirection.Up:
                return Vector3.up;
                break;
            case SampleDirection.Down:
                return Vector3.down;
                break;
            case SampleDirection.Left:
                return Vector3.left;
                break;
            case SampleDirection.Right:
                return Vector3.right;
                break;
        }
        
        return Vector3.zero;
    }
    
    public SphericalHarmonicsL2 DecodeSH(CompressedProbeData d, int idx)
    {
        var sh = new SphericalHarmonicsL2();
        // int k = d.coefficients[idx].Length;
        int k = d.coefficients.ColCount;
        for (int rgb = 0; rgb < 3; rgb++)
            for (int j = 0; j < 9; j++)
            {
                float sum = 0f;
                for (int c = 0; c < k; c++)
                    sum += d.coefficients[idx, c] * d.pcaAxes[c, rgb * 9 + j];
                sh[rgb, j] = sum + d.mean[rgb * 9 + j];
            }
        return sh;
    }
    
    public static string SphericalHarmonicsToString(SphericalHarmonicsL2 sh)
    {
        StringBuilder sb = new StringBuilder();
        for (int channel = 0; channel < 3; channel++) // R, G, B
        {
            sb.AppendLine($"Channel {(char)('R' + channel)}:");
            for (int i = 0; i < 9; i++)
            {
                float coefficient = sh[channel, i];
                sb.AppendLine($"  L{i}: {coefficient:F6}");
            }
        }
        return sb.ToString();
    }
}
